{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "284d3de2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !sudo apt install libopenmpi-dev -y\n",
    "# !pip3 install mpi4py --user\n",
    "# !pip3 install deepspeed==0.12.3 --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e4bc859d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip3 install accelerate transformers -U --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "992d1de9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absl-py==2.0.0\r\n",
      "accelerate==0.25.0\r\n",
      "aiofiles==23.2.1\r\n",
      "aiohttp==3.8.5\r\n",
      "aiohttp-cors==0.7.0\r\n",
      "aiorwlock==1.3.0\r\n",
      "aiosignal==1.3.1\r\n",
      "altair==5.1.2\r\n",
      "anyio==3.7.1\r\n",
      "appdirs==1.4.4\r\n",
      "argon2-cffi==23.1.0\r\n",
      "argon2-cffi-bindings==21.2.0\r\n",
      "asttokens==2.2.1\r\n",
      "async-timeout==4.0.3\r\n",
      "attributedict==0.3.0\r\n",
      "attrs==23.1.0\r\n",
      "autoawq @ https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.6/autoawq-0.1.6+cu118-cp310-cp310-linux_x86_64.whl\r\n",
      "azure-core==1.29.5\r\n",
      "azure-identity==1.15.0\r\n",
      "azure-storage-blob==12.18.3\r\n",
      "azure-storage-file-datalake==12.13.2\r\n",
      "backcall==0.2.0\r\n",
      "bcrypt==4.0.1\r\n",
      "beautifulsoup4==4.12.2\r\n",
      "bitsandbytes==0.41.0\r\n",
      "bleach==6.0.0\r\n",
      "blessed==1.20.0\r\n",
      "blessings==1.7\r\n",
      "boto3==1.28.78\r\n",
      "botocore==1.31.78\r\n",
      "Brotli==1.1.0\r\n",
      "cachetools==5.3.2\r\n",
      "causal-conv1d==1.0.0\r\n",
      "certifi==2022.12.7\r\n",
      "cffi==1.15.1\r\n",
      "chardet==5.2.0\r\n",
      "charset-normalizer==2.1.1\r\n",
      "circuitbreaker==1.4.0\r\n",
      "click==8.1.7\r\n",
      "cmake==3.27.7\r\n",
      "codecov==2.1.13\r\n",
      "colorama==0.4.6\r\n",
      "coloredlogs==15.0.1\r\n",
      "colorful==0.5.5\r\n",
      "colour-runner==0.1.1\r\n",
      "comm==0.1.4\r\n",
      "contourpy==1.2.0\r\n",
      "coverage==7.3.2\r\n",
      "cryptography==41.0.5\r\n",
      "cycler==0.12.1\r\n",
      "DataProperty==1.0.1\r\n",
      "datasets==2.14.6\r\n",
      "debugpy==1.6.7.post1\r\n",
      "decorator==5.1.1\r\n",
      "deepdiff==6.6.1\r\n",
      "deepspeed==0.12.3\r\n",
      "defusedxml==0.7.1\r\n",
      "dill==0.3.7\r\n",
      "distlib==0.3.7\r\n",
      "distro==1.7.0\r\n",
      "docker-pycreds==0.4.0\r\n",
      "einops==0.6.1\r\n",
      "einops-exts==0.0.4\r\n",
      "evaluate==0.4.1\r\n",
      "exceptiongroup==1.1.3\r\n",
      "executing==1.2.0\r\n",
      "fastapi==0.104.1\r\n",
      "fastjsonschema==2.18.0\r\n",
      "ffmpy==0.3.1\r\n",
      "filelock==3.13.1\r\n",
      "flash-attn==2.3.0\r\n",
      "fonttools==4.44.0\r\n",
      "frozenlist==1.4.0\r\n",
      "fsspec==2023.10.0\r\n",
      "gitdb==4.0.11\r\n",
      "GitPython==3.1.40\r\n",
      "google-api-core==2.12.0\r\n",
      "google-auth==2.23.4\r\n",
      "google-cloud-core==2.3.3\r\n",
      "google-cloud-storage==2.10.0\r\n",
      "google-crc32c==1.5.0\r\n",
      "google-resumable-media==2.6.0\r\n",
      "googleapis-common-protos==1.61.0\r\n",
      "gpustat==1.1.1\r\n",
      "gradio==3.35.2\r\n",
      "gradio_client==0.2.9\r\n",
      "grpcio==1.59.2\r\n",
      "h11==0.14.0\r\n",
      "hjson==3.1.0\r\n",
      "httpcore==0.17.3\r\n",
      "httptools==0.6.1\r\n",
      "httpx==0.24.0\r\n",
      "huggingface-hub==0.17.3\r\n",
      "humanfriendly==10.0\r\n",
      "idna==3.4\r\n",
      "inspecta==0.1.3\r\n",
      "ipykernel==6.25.1\r\n",
      "ipython==8.14.0\r\n",
      "ipython-genutils==0.2.0\r\n",
      "ipywidgets==8.1.0\r\n",
      "isodate==0.6.1\r\n",
      "jedi==0.19.0\r\n",
      "Jinja2==3.1.2\r\n",
      "jmespath==1.0.1\r\n",
      "joblib==1.3.2\r\n",
      "jsonlines==4.0.0\r\n",
      "jsonschema==4.19.0\r\n",
      "jsonschema-specifications==2023.7.1\r\n",
      "jupyter==1.0.0\r\n",
      "jupyter-console==6.6.3\r\n",
      "jupyter-server==1.18.0\r\n",
      "jupyter-server-proxy==3.2.1\r\n",
      "jupyter_client==8.3.0\r\n",
      "jupyter_core==5.3.1\r\n",
      "jupyterlab-pygments==0.2.2\r\n",
      "jupyterlab-widgets==3.0.8\r\n",
      "kiwisolver==1.4.5\r\n",
      "linkify-it-py==2.0.2\r\n",
      "lit==17.0.4\r\n",
      "# Editable install with no version control (llava==1.1.3)\r\n",
      "-e /home/ubuntu/LLaVA\r\n",
      "lm-eval==0.3.0\r\n",
      "mamba-ssm @ file:///home/ubuntu/mamba\r\n",
      "markdown-it-py==2.2.0\r\n",
      "markdown2==2.4.10\r\n",
      "MarkupSafe==2.1.3\r\n",
      "matplotlib==3.8.1\r\n",
      "matplotlib-inline==0.1.6\r\n",
      "mbstrdecoder==1.1.3\r\n",
      "mdit-py-plugins==0.3.3\r\n",
      "mdurl==0.1.2\r\n",
      "mistune==3.0.1\r\n",
      "mosaicml-streaming==0.6.1\r\n",
      "mp==0.5.0\r\n",
      "mpi4py==3.1.5\r\n",
      "mpmath==1.3.0\r\n",
      "msal==1.25.0\r\n",
      "msal-extensions==1.0.0\r\n",
      "msgpack==1.0.7\r\n",
      "msgspec==0.18.4\r\n",
      "multidict==6.0.4\r\n",
      "multiprocess==0.70.15\r\n",
      "nbclient==0.8.0\r\n",
      "nbconvert==7.7.4\r\n",
      "nbformat==5.9.2\r\n",
      "nest-asyncio==1.5.7\r\n",
      "networkx==3.0\r\n",
      "ninja==1.11.1.1\r\n",
      "nltk==3.8.1\r\n",
      "notebook==6.4.12\r\n",
      "numexpr==2.8.7\r\n",
      "numpy==1.24.1\r\n",
      "nvidia-cublas-cu11==11.10.3.66\r\n",
      "nvidia-cuda-cupti-cu11==11.7.101\r\n",
      "nvidia-cuda-nvrtc-cu11==11.7.99\r\n",
      "nvidia-cuda-runtime-cu11==11.7.99\r\n",
      "nvidia-cudnn-cu11==8.5.0.96\r\n",
      "nvidia-cufft-cu11==10.9.0.58\r\n",
      "nvidia-curand-cu11==10.2.10.91\r\n",
      "nvidia-cusolver-cu11==11.4.0.1\r\n",
      "nvidia-cusparse-cu11==11.7.4.91\r\n",
      "nvidia-ml-py==12.535.133\r\n",
      "nvidia-nccl-cu11==2.14.3\r\n",
      "nvidia-nvtx-cu11==11.7.91\r\n",
      "oci==2.115.0\r\n",
      "openai==0.28.0\r\n",
      "opencensus==0.11.3\r\n",
      "opencensus-context==0.1.3\r\n",
      "ordered-set==4.1.0\r\n",
      "orjson==3.9.10\r\n",
      "packaging==23.1\r\n",
      "pandas==2.1.2\r\n",
      "pandocfilters==1.5.0\r\n",
      "paramiko==3.3.1\r\n",
      "parso==0.8.3\r\n",
      "pathvalidate==3.2.0\r\n",
      "peft==0.4.0\r\n",
      "pexpect==4.8.0\r\n",
      "pickleshare==0.7.5\r\n",
      "Pillow==9.3.0\r\n",
      "platformdirs==3.10.0\r\n",
      "pluggy==1.3.0\r\n",
      "portalocker==2.8.2\r\n",
      "prometheus-client==0.17.1\r\n",
      "prompt-toolkit==3.0.39\r\n",
      "protobuf==4.25.0\r\n",
      "psutil==5.9.5\r\n",
      "ptyprocess==0.7.0\r\n",
      "pure-eval==0.2.2\r\n",
      "py-cpuinfo==9.0.0\r\n",
      "py-spy==0.3.14\r\n",
      "pyarrow==14.0.0\r\n",
      "pyasn1==0.5.0\r\n",
      "pyasn1-modules==0.3.0\r\n",
      "pybind11==2.11.1\r\n",
      "pycountry==22.3.5\r\n",
      "pycparser==2.21\r\n",
      "pydantic==1.10.13\r\n",
      "pydub==0.25.1\r\n",
      "Pygments==2.16.1\r\n",
      "PyJWT==2.8.0\r\n",
      "PyNaCl==1.5.0\r\n",
      "pynvml==11.5.0\r\n",
      "pyOpenSSL==23.3.0\r\n",
      "pyparsing==3.1.1\r\n",
      "pyproject-api==1.6.1\r\n",
      "pytablewriter==1.2.0\r\n",
      "python-dateutil==2.8.2\r\n",
      "python-dotenv==1.0.0\r\n",
      "python-multipart==0.0.6\r\n",
      "python-snappy==0.6.1\r\n",
      "pytz==2023.3.post1\r\n",
      "PyYAML==6.0.1\r\n",
      "pyzmq==25.1.1\r\n",
      "qtconsole==5.4.3\r\n",
      "QtPy==2.3.1\r\n",
      "ray==2.8.0\r\n",
      "referencing==0.30.2\r\n",
      "regex==2023.10.3\r\n",
      "requests==2.28.1\r\n",
      "responses==0.18.0\r\n",
      "rich==13.6.0\r\n",
      "rootpath==0.1.1\r\n",
      "rouge-score==0.1.2\r\n",
      "rpds-py==0.9.2\r\n",
      "rsa==4.9\r\n",
      "s3transfer==0.7.0\r\n",
      "sacrebleu==1.5.0\r\n",
      "safetensors==0.4.0\r\n",
      "scikit-learn==1.2.2\r\n",
      "scipy==1.11.3\r\n",
      "semantic-version==2.10.0\r\n",
      "Send2Trash==1.8.2\r\n",
      "sentencepiece==0.1.99\r\n",
      "sentry-sdk==1.36.0\r\n",
      "setproctitle==1.3.3\r\n",
      "shortuuid==1.0.11\r\n",
      "simpervisor==1.0.0\r\n",
      "six==1.16.0\r\n",
      "smart-open==6.4.0\r\n",
      "smmap==5.0.1\r\n",
      "sniffio==1.3.0\r\n",
      "soupsieve==2.4.1\r\n",
      "sqlitedict==2.1.0\r\n",
      "ssh-import-id==5.11\r\n",
      "stack-data==0.6.2\r\n",
      "starlette==0.27.0\r\n",
      "svgwrite==1.4.3\r\n",
      "sympy==1.12\r\n",
      "tabledata==1.3.3\r\n",
      "tabulate==0.9.0\r\n",
      "tcolorpy==0.1.4\r\n",
      "tensorboardX==2.6.2.2\r\n",
      "termcolor==2.3.0\r\n",
      "terminado==0.17.1\r\n",
      "texttable==1.7.0\r\n",
      "threadpoolctl==3.2.0\r\n",
      "timm==0.6.13\r\n",
      "tinycss2==1.2.1\r\n",
      "tokenizers==0.14.1\r\n",
      "toml==0.10.2\r\n",
      "tomli==2.0.1\r\n",
      "toolz==0.12.0\r\n",
      "torch==2.0.1+cu118\r\n",
      "torchaudio==2.0.2+cu118\r\n",
      "torchvision==0.15.2+cu118\r\n",
      "tornado==6.3.3\r\n",
      "tox==4.11.3\r\n",
      "tqdm==4.66.1\r\n",
      "tqdm-multiprocess==0.0.11\r\n",
      "traitlets==5.9.0\r\n",
      "transformers==4.35.2\r\n",
      "triton==2.0.0\r\n",
      "typepy==1.3.2\r\n",
      "typing_extensions==4.8.0\r\n",
      "tzdata==2023.3\r\n",
      "uc-micro-py==1.0.2\r\n",
      "Unidecode==1.3.7\r\n",
      "unzip==1.0.0\r\n",
      "urllib3==1.26.13\r\n",
      "uvicorn==0.24.0.post1\r\n",
      "uvloop==0.19.0\r\n",
      "virtualenv==20.21.0\r\n",
      "wandb==0.16.0\r\n",
      "watchfiles==0.21.0\r\n",
      "wavedrom==2.0.3.post3\r\n",
      "wcwidth==0.2.6\r\n",
      "webencodings==0.5.1\r\n",
      "websocket-client==1.6.1\r\n",
      "websockets==12.0\r\n",
      "widgetsnbextension==4.0.8\r\n",
      "xxhash==3.4.1\r\n",
      "yarl==1.9.2\r\n",
      "zstandard==0.22.0\r\n",
      "zstd==1.5.5.1\r\n"
     ]
    }
   ],
   "source": [
    "!pip3 freeze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ecf642c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tue Dec  5 16:04:56 2023       \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |\r\n",
      "|-------------------------------+----------------------+----------------------+\r\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\r\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\r\n",
      "|                               |                      |               MIG M. |\r\n",
      "|===============================+======================+======================|\r\n",
      "|   0  NVIDIA A100 80G...  On   | 00000001:00:00.0 Off |                    0 |\r\n",
      "| N/A   39C    P0    63W / 300W |  13456MiB / 81920MiB |      0%      Default |\r\n",
      "|                               |                      |             Disabled |\r\n",
      "+-------------------------------+----------------------+----------------------+\r\n",
      "                                                                               \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| Processes:                                                                  |\r\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\r\n",
      "|        ID   ID                                                   Usage      |\r\n",
      "|=============================================================================|\r\n",
      "+-----------------------------------------------------------------------------+\r\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b2374da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9b310aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2023, Albert Gu, Tri Dao.\n",
    "\n",
    "import math\n",
    "from functools import partial\n",
    "\n",
    "from collections import namedtuple\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import CrossEntropyLoss\n",
    "from transformers import PretrainedConfig, PreTrainedModel\n",
    "\n",
    "from mamba_ssm.modules.mamba_simple import Mamba, Block\n",
    "from mamba_ssm.utils.generation import GenerationMixin\n",
    "from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf\n",
    "\n",
    "try:\n",
    "    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn\n",
    "except ImportError:\n",
    "    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None\n",
    "\n",
    "\n",
    "def create_block(\n",
    "    d_model,\n",
    "    ssm_cfg=None,\n",
    "    norm_epsilon=1e-5,\n",
    "    rms_norm=False,\n",
    "    residual_in_fp32=False,\n",
    "    fused_add_norm=False,\n",
    "    layer_idx=None,\n",
    "    device=None,\n",
    "    dtype=None,\n",
    "):\n",
    "    if ssm_cfg is None:\n",
    "        ssm_cfg = {}\n",
    "    factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
    "    mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)\n",
    "    norm_cls = partial(\n",
    "        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs\n",
    "    )\n",
    "    block = Block(\n",
    "        d_model,\n",
    "        mixer_cls,\n",
    "        norm_cls=norm_cls,\n",
    "        fused_add_norm=fused_add_norm,\n",
    "        residual_in_fp32=residual_in_fp32,\n",
    "    )\n",
    "    block.layer_idx = layer_idx\n",
    "    return block\n",
    "\n",
    "\n",
    "# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454\n",
    "def _init_weights(\n",
    "    module,\n",
    "    n_layer,\n",
    "    initializer_range=0.02,  # Now only used for embedding layer.\n",
    "    rescale_prenorm_residual=True,\n",
    "    n_residuals_per_layer=1,  # Change to 2 if we have MLP\n",
    "):\n",
    "    if isinstance(module, nn.Linear):\n",
    "        if module.bias is not None:\n",
    "            if not getattr(module.bias, \"_no_reinit\", False):\n",
    "                nn.init.zeros_(module.bias)\n",
    "    elif isinstance(module, nn.Embedding):\n",
    "        nn.init.normal_(module.weight, std=initializer_range)\n",
    "\n",
    "    if rescale_prenorm_residual:\n",
    "        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:\n",
    "        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale\n",
    "        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.\n",
    "        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/\n",
    "        #\n",
    "        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py\n",
    "        for name, p in module.named_parameters():\n",
    "            if name in [\"out_proj.weight\", \"fc2.weight\"]:\n",
    "                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block\n",
    "                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)\n",
    "                # We need to reinit p since this code could be called multiple times\n",
    "                # Having just p *= scale would repeatedly scale it down\n",
    "                nn.init.kaiming_uniform_(p, a=math.sqrt(5))\n",
    "                with torch.no_grad():\n",
    "                    p /= math.sqrt(n_residuals_per_layer * n_layer)\n",
    "\n",
    "\n",
    "class MixerModel(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        d_model: int,\n",
    "        n_layer: int,\n",
    "        vocab_size: int,\n",
    "        ssm_cfg=None,\n",
    "        norm_epsilon: float = 1e-5,\n",
    "        rms_norm: bool = False,\n",
    "        initializer_cfg=None,\n",
    "        fused_add_norm=False,\n",
    "        residual_in_fp32=False,\n",
    "        device=None,\n",
    "        dtype=None,\n",
    "    ) -> None:\n",
    "        factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
    "        super().__init__()\n",
    "        self.residual_in_fp32 = residual_in_fp32\n",
    "\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)\n",
    "\n",
    "        # We change the order of residual and layer norm:\n",
    "        # Instead of LN -> Attn / MLP -> Add, we do:\n",
    "        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and\n",
    "        # the main branch (output of MLP / Mixer). The model definition is unchanged.\n",
    "        # This is for performance reason: we can fuse add + layer_norm.\n",
    "        self.fused_add_norm = fused_add_norm\n",
    "        if self.fused_add_norm:\n",
    "            if layer_norm_fn is None or rms_norm_fn is None:\n",
    "                raise ImportError(\"Failed to import Triton LayerNorm / RMSNorm kernels\")\n",
    "\n",
    "        self.layers = nn.ModuleList(\n",
    "            [\n",
    "                create_block(\n",
    "                    d_model,\n",
    "                    ssm_cfg=ssm_cfg,\n",
    "                    norm_epsilon=norm_epsilon,\n",
    "                    rms_norm=rms_norm,\n",
    "                    residual_in_fp32=residual_in_fp32,\n",
    "                    fused_add_norm=fused_add_norm,\n",
    "                    layer_idx=i,\n",
    "                    **factory_kwargs,\n",
    "                )\n",
    "                for i in range(n_layer)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(\n",
    "            d_model, eps=norm_epsilon, **factory_kwargs\n",
    "        )\n",
    "\n",
    "        self.apply(\n",
    "            partial(\n",
    "                _init_weights,\n",
    "                n_layer=n_layer,\n",
    "                **(initializer_cfg if initializer_cfg is not None else {}),\n",
    "            )\n",
    "        )\n",
    "\n",
    "    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):\n",
    "        return {\n",
    "            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)\n",
    "            for i, layer in enumerate(self.layers)\n",
    "        }\n",
    "\n",
    "    def forward(self, input_ids, inference_params=None):\n",
    "        hidden_states = self.embedding(input_ids)\n",
    "        residual = None\n",
    "        for layer in self.layers:\n",
    "            hidden_states, residual = layer(\n",
    "                hidden_states, residual, inference_params=inference_params\n",
    "            )\n",
    "        if not self.fused_add_norm:\n",
    "            residual = (hidden_states + residual) if residual is not None else hidden_states\n",
    "            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))\n",
    "        else:\n",
    "            # Set prenorm=False here since we don't need the residual\n",
    "            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn\n",
    "            hidden_states = fused_add_norm_fn(\n",
    "                hidden_states,\n",
    "                self.norm_f.weight,\n",
    "                self.norm_f.bias,\n",
    "                eps=self.norm_f.eps,\n",
    "                residual=residual,\n",
    "                prenorm=False,\n",
    "                residual_in_fp32=self.residual_in_fp32,\n",
    "            )\n",
    "        return hidden_states\n",
    "\n",
    "class MambaLMHeadModel(PreTrainedModel, GenerationMixin):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        config,\n",
    "        initializer_cfg=None,\n",
    "        pad_vocab_size_multiple: int = 1,\n",
    "        device=None,\n",
    "        dtype=None,\n",
    "        **backbone_kwargs,\n",
    "    ) -> None:\n",
    "        d_model = config.d_model\n",
    "        n_layer = config.n_layer\n",
    "        vocab_size = config.vocab_size\n",
    "        factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
    "        super().__init__(config=config)\n",
    "        if vocab_size % pad_vocab_size_multiple != 0:\n",
    "            vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)\n",
    "        self.backbone = MixerModel(\n",
    "            d_model=d_model,\n",
    "            n_layer=n_layer,\n",
    "            vocab_size=vocab_size,\n",
    "            initializer_cfg=initializer_cfg,\n",
    "            **backbone_kwargs,\n",
    "            **factory_kwargs,\n",
    "        )\n",
    "        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)\n",
    "\n",
    "        # Initialize weights and apply final processing\n",
    "        self.apply(\n",
    "            partial(\n",
    "                _init_weights,\n",
    "                n_layer=n_layer,\n",
    "                **(initializer_cfg if initializer_cfg is not None else {}),\n",
    "            )\n",
    "        )\n",
    "        self.tie_weights()\n",
    "        # _tied_weights_keys = ['lm_head.weight']\n",
    "\n",
    "    def tie_weights(self):\n",
    "        self.lm_head.weight = self.backbone.embedding.weight\n",
    "\n",
    "    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):\n",
    "        return self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)\n",
    "\n",
    "    def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0, labels = None):\n",
    "        \"\"\"\n",
    "        \"position_ids\" is just to be compatible with Transformer generation. We don't use it.\n",
    "        num_last_tokens: if > 0, only return the logits for the last n tokens\n",
    "        \"\"\"\n",
    "        hidden_states = self.backbone(input_ids, inference_params=inference_params)\n",
    "        if num_last_tokens > 0:\n",
    "            hidden_states = hidden_states[:, -num_last_tokens:]\n",
    "        lm_logits = self.lm_head(hidden_states)\n",
    "        \n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            logits = lm_logits\n",
    "            # Shift so that tokens < n predict n\n",
    "            shift_logits = logits[..., :-1, :].contiguous()\n",
    "            shift_labels = labels[..., 1:].contiguous()\n",
    "            # Flatten the tokens\n",
    "            loss_fct = CrossEntropyLoss()\n",
    "            shift_logits = shift_logits.view(-1, self.config.vocab_size)\n",
    "            shift_labels = shift_labels.view(-1)\n",
    "            # Enable model parallelism\n",
    "            shift_labels = shift_labels.to(shift_logits.device)\n",
    "            loss = loss_fct(shift_logits, shift_labels)\n",
    "            print(loss, shift_logits, shift_logits.dtype, shift_labels, shift_labels.dtype)\n",
    "            return (loss,)\n",
    "            \n",
    "        else:\n",
    "            CausalLMOutput = namedtuple(\"CausalLMOutput\", [\"logits\"])\n",
    "            return CausalLMOutput(logits=lm_logits)\n",
    "\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs):\n",
    "        config = load_config_hf(pretrained_model_name)\n",
    "        model = cls(**config, device=device, dtype=dtype, **kwargs)\n",
    "        model.load_state_dict(load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype))\n",
    "        return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4fe72f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://huggingface.co/state-spaces/mamba-1.4b/raw/main/config.json -O config-1.4b.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eac72d2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('config-1.4b.json') as fopen:\n",
    "    config = json.load(fopen)\n",
    "    config['hidden_size'] = config['d_model']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "d17e4c36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "PretrainedConfig {\n",
       "  \"d_model\": 2048,\n",
       "  \"fused_add_norm\": true,\n",
       "  \"hidden_size\": 2048,\n",
       "  \"n_layer\": 48,\n",
       "  \"pad_vocab_size_multiple\": 8,\n",
       "  \"residual_in_fp32\": true,\n",
       "  \"rms_norm\": true,\n",
       "  \"ssm_cfg\": {},\n",
       "  \"transformers_version\": \"4.35.2\",\n",
       "  \"vocab_size\": 32000\n",
       "}"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "config = PretrainedConfig(**{**config, 'vocab_size': 32000})\n",
    "config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b85b48be",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MambaLMHeadModel(config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a3daef27",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MambaLMHeadModel(\n",
       "  (backbone): MixerModel(\n",
       "    (embedding): Embedding(32000, 2048)\n",
       "    (layers): ModuleList(\n",
       "      (0-47): 48 x Block(\n",
       "        (mixer): Mamba(\n",
       "          (in_proj): Linear(in_features=2048, out_features=8192, bias=False)\n",
       "          (conv1d): Conv1d(4096, 4096, kernel_size=(4,), stride=(1,), padding=(3,), groups=4096)\n",
       "          (act): SiLU()\n",
       "          (x_proj): Linear(in_features=4096, out_features=160, bias=False)\n",
       "          (dt_proj): Linear(in_features=128, out_features=4096, bias=True)\n",
       "          (out_proj): Linear(in_features=4096, out_features=2048, bias=False)\n",
       "        )\n",
       "        (norm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "    )\n",
       "    (norm_f): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=2048, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e8beea8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from streaming import LocalDataset\n",
    "import numpy as np\n",
    "from streaming.base.format.mds.encodings import Encoding, _encodings\n",
    "\n",
    "class UInt16(Encoding):\n",
    "    def encode(self, obj) -> bytes:\n",
    "        return obj.tobytes()\n",
    "\n",
    "    def decode(self, data: bytes):\n",
    "        return np.frombuffer(data, np.uint16)\n",
    "\n",
    "_encodings['uint16'] = UInt16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "74c0f1fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !git lfs clone https://huggingface.co/datasets/malaysia-ai/mosaic-instructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "88f14bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatasetFixed(torch.utils.data.Dataset):\n",
    "    def __init__(self, local):\n",
    "        self.dataset = LocalDataset(local=local)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        print(idx)\n",
    "        data = self.dataset[idx]\n",
    "        data['labels'] = data['input_ids'].copy()\n",
    "\n",
    "        data.pop('token_type_ids', None)\n",
    "        for k in data.keys():\n",
    "            data[k] = data[k].astype(np.int64)\n",
    "        return data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "train_dataset = DatasetFixed(local='mosaic-instructions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "7ea1956b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n"
     ]
    }
   ],
   "source": [
    "from transformers import TrainingArguments, Trainer, default_data_collator\n",
    "\n",
    "output_dir = 'test-1.4b'\n",
    "\n",
    "deepspeed = {\n",
    "    \"comms_logger\": {\n",
    "        \"enabled\": True,\n",
    "        \"debug\": True\n",
    "    },\n",
    "    \"fp16\": {\n",
    "        \"enabled\": \"auto\",\n",
    "        \"loss_scale\": 0,\n",
    "        \"loss_scale_window\": 1000,\n",
    "        \"initial_scale_power\": 16,\n",
    "        \"hysteresis\": 2,\n",
    "        \"min_loss_scale\": 1\n",
    "    },\n",
    "\n",
    "    \"bf16\": {\n",
    "        \"enabled\": \"auto\"\n",
    "    },\n",
    "\n",
    "    \"optimizer\": {\n",
    "        \"type\": \"AdamW\",\n",
    "        \"params\": {\n",
    "            \"lr\": \"auto\",\n",
    "            \"betas\": \"auto\",\n",
    "            \"eps\": \"auto\",\n",
    "            \"weight_decay\": \"auto\"\n",
    "        }\n",
    "    },\n",
    "\n",
    "    \"scheduler\": {\n",
    "        \"type\": \"WarmupDecayLR\",\n",
    "        \"params\": {\n",
    "            \"warmup_min_lr\": \"auto\",\n",
    "            \"warmup_max_lr\": \"auto\",\n",
    "            \"warmup_num_steps\": \"auto\",\n",
    "            \"total_num_steps\": \"auto\",\n",
    "        }\n",
    "    },\n",
    "\n",
    "    \"zero_optimization\": {\n",
    "        \"stage\": 3,\n",
    "        \"offload_optimizer\": {\n",
    "            \"device\": \"cpu\",\n",
    "            \"pin_memory\": True\n",
    "        },\n",
    "        \"offload_param\": {\n",
    "            \"device\": \"cpu\",\n",
    "            \"pin_memory\": True\n",
    "        },\n",
    "        \"overlap_comm\": True,\n",
    "        \"contiguous_gradients\": True,\n",
    "        \"sub_group_size\": 1e8,\n",
    "        \"reduce_bucket_size\": \"auto\",\n",
    "        \"stage3_prefetch_bucket_size\": \"auto\",\n",
    "        \"stage3_param_persistence_threshold\": \"auto\",\n",
    "        \"stage3_max_live_parameters\": 1e8,\n",
    "        \"stage3_max_reuse_distance\": 1e8,\n",
    "        \"stage3_gather_16bit_weights_on_model_save\": True\n",
    "    },\n",
    "\n",
    "    \"gradient_accumulation_steps\": \"auto\",\n",
    "    \"gradient_clipping\": \"auto\",\n",
    "    \"steps_per_print\": 2000,\n",
    "    \"train_batch_size\": \"auto\",\n",
    "    \"train_micro_batch_size_per_gpu\": \"auto\",\n",
    "    \"wall_clock_breakdown\": False\n",
    "}\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir,\n",
    "    per_device_train_batch_size=2,\n",
    "    gradient_accumulation_steps=1,\n",
    "    logging_steps=1,\n",
    "    save_strategy='steps',\n",
    "    save_steps=5,\n",
    "    num_train_epochs=3,\n",
    "    learning_rate=1e-4,\n",
    "    weight_decay=0,\n",
    "    warmup_steps=1000,\n",
    "    bf16=True,\n",
    "    fp16=False,\n",
    "    gradient_checkpointing=False,\n",
    "    deepspeed=deepspeed,\n",
    "    save_total_limit=5,\n",
    "    log_level='debug',\n",
    "    max_steps=100,\n",
    "    save_safetensors=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "d3e45910",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "max_steps is given, it will override any value given in num_train_epochs\n",
      "Using auto half precision backend\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    data_collator=default_data_collator,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "edabeb16",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "MambaLMHeadModel(\n",
       "  (backbone): MixerModel(\n",
       "    (embedding): Embedding(32000, 2048)\n",
       "    (layers): ModuleList(\n",
       "      (0-47): 48 x Block(\n",
       "        (mixer): Mamba(\n",
       "          (in_proj): Linear(in_features=2048, out_features=8192, bias=False)\n",
       "          (conv1d): Conv1d(4096, 4096, kernel_size=(4,), stride=(1,), padding=(3,), groups=4096)\n",
       "          (act): SiLU()\n",
       "          (x_proj): Linear(in_features=4096, out_features=160, bias=False)\n",
       "          (dt_proj): Linear(in_features=128, out_features=4096, bias=True)\n",
       "          (out_proj): Linear(in_features=4096, out_features=2048, bias=False)\n",
       "        )\n",
       "        (norm): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "      )\n",
       "    )\n",
       "    (norm_f): LayerNorm((2048,), eps=1e-05, elementwise_affine=True)\n",
       "  )\n",
       "  (lm_head): Linear(in_features=2048, out_features=32000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "5526e37b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# trainer.train()\n",
    "# already saved 2 checkpoints, now want to test to load"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1c6775c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Currently training with a batch size of: 2\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-12-05 16:05:11,549] [INFO] [real_accelerator.py:158:get_accelerator] Setting ds_accelerator to cuda (auto detect)\n",
      "[2023-12-05 16:05:13,847] [INFO] [logging.py:96:log_dist] [Rank -1] DeepSpeed info: version=0.12.3, git-hash=unknown, git-branch=unknown\n",
      "[2023-12-05 16:05:13,848] [INFO] [comm.py:637:init_distributed] cdb=None\n",
      "[2023-12-05 16:05:13,848] [INFO] [comm.py:652:init_distributed] Not using the DeepSpeed or dist launchers, attempting to detect MPI environment...\n",
      "[2023-12-05 16:05:14,243] [INFO] [comm.py:702:mpi_discovery] Discovered MPI settings of world_rank=0, local_rank=0, world_size=1, master_addr=10.208.0.238, master_port=29500\n",
      "[2023-12-05 16:05:14,244] [INFO] [comm.py:668:init_distributed] Initializing TorchBackend in DeepSpeed with backend nccl\n",
      "[2023-12-05 16:05:16,230] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using /home/ubuntu/.cache/torch_extensions/py310_cu118 as PyTorch extensions root...\n",
      "Detected CUDA files, patching ldflags\n",
      "Emitting ninja build file /home/ubuntu/.cache/torch_extensions/py310_cu118/cpu_adam/build.ninja...\n",
      "Building extension module cpu_adam...\n",
      "Allowing ninja to set a default number of workers... (overridable by setting the environment variable MAX_JOBS=N)\n",
      "Loading extension module cpu_adam...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ninja: no work to do.\n",
      "Time to load cpu_adam op: 2.2710254192352295 seconds\n",
      "[2023-12-05 16:05:20,706] [INFO] [logging.py:96:log_dist] [Rank 0] Using DeepSpeed Optimizer param name adamw as basic optimizer\n",
      "[2023-12-05 16:05:20,706] [INFO] [logging.py:96:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer\n",
      "[2023-12-05 16:05:20,735] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Basic Optimizer = DeepSpeedCPUAdam\n",
      "[2023-12-05 16:05:20,735] [INFO] [utils.py:56:is_zero_supported_optimizer] Checking ZeRO support for optimizer=DeepSpeedCPUAdam type=<class 'deepspeed.ops.adam.cpu_adam.DeepSpeedCPUAdam'>\n",
      "[2023-12-05 16:05:20,735] [INFO] [logging.py:96:log_dist] [Rank 0] Creating fp16 ZeRO stage 3 optimizer, MiCS is enabled False, Hierarchical params gather False\n",
      "[2023-12-05 16:05:20,736] [INFO] [logging.py:96:log_dist] [Rank 0] Creating torch.bfloat16 ZeRO stage 3 optimizer\n",
      "[2023-12-05 16:05:20,856] [INFO] [utils.py:802:see_memory_usage] Stage 3 initialize beginning\n",
      "[2023-12-05 16:05:20,857] [INFO] [utils.py:803:see_memory_usage] MA 2.49 GB         Max_MA 2.49 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:20,858] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 23.28 GB, percent = 10.8%\n",
      "[2023-12-05 16:05:20,860] [INFO] [stage3.py:127:__init__] Reduce bucket size 4194304\n",
      "[2023-12-05 16:05:20,860] [INFO] [stage3.py:128:__init__] Prefetch bucket size 3774873\n",
      "[2023-12-05 16:05:20,972] [INFO] [utils.py:802:see_memory_usage] DeepSpeedZeRoOffload initialize [begin]\n",
      "[2023-12-05 16:05:20,973] [INFO] [utils.py:803:see_memory_usage] MA 2.49 GB         Max_MA 2.49 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:20,974] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 23.28 GB, percent = 10.8%\n",
      "Parameter Offload: Total persistent parameters: 1576960 in 290 params\n",
      "[2023-12-05 16:05:22,465] [INFO] [utils.py:802:see_memory_usage] DeepSpeedZeRoOffload initialize [end]\n",
      "[2023-12-05 16:05:22,466] [INFO] [utils.py:803:see_memory_usage] MA 0.0 GB         Max_MA 2.49 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:22,467] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 25.92 GB, percent = 12.0%\n",
      "[2023-12-05 16:05:22,585] [INFO] [utils.py:802:see_memory_usage] Before creating fp16 partitions\n",
      "[2023-12-05 16:05:22,586] [INFO] [utils.py:803:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:22,587] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 25.92 GB, percent = 12.0%\n",
      "[2023-12-05 16:05:24,039] [INFO] [utils.py:802:see_memory_usage] After creating fp16 partitions: 13\n",
      "[2023-12-05 16:05:24,040] [INFO] [utils.py:803:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:24,041] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 29.99 GB, percent = 13.9%\n",
      "[2023-12-05 16:05:24,157] [INFO] [utils.py:802:see_memory_usage] Before creating fp32 partitions\n",
      "[2023-12-05 16:05:24,158] [INFO] [utils.py:803:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:24,159] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 29.99 GB, percent = 13.9%\n",
      "[2023-12-05 16:05:24,636] [INFO] [utils.py:802:see_memory_usage] After creating fp32 partitions\n",
      "[2023-12-05 16:05:24,638] [INFO] [utils.py:803:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:24,638] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 34.94 GB, percent = 16.2%\n",
      "[2023-12-05 16:05:24,756] [INFO] [utils.py:802:see_memory_usage] Before initializing optimizer states\n",
      "[2023-12-05 16:05:24,757] [INFO] [utils.py:803:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:24,757] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 34.94 GB, percent = 16.2%\n",
      "[2023-12-05 16:05:27,862] [INFO] [utils.py:802:see_memory_usage] After initializing optimizer states\n",
      "[2023-12-05 16:05:27,863] [INFO] [utils.py:803:see_memory_usage] MA 0.0 GB         Max_MA 0.0 GB         CA 2.49 GB         Max_CA 2 GB \n",
      "[2023-12-05 16:05:27,864] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 51.28 GB, percent = 23.7%\n",
      "[2023-12-05 16:05:27,865] [INFO] [stage3.py:479:_setup_for_real_optimizer] optimizer state initialized\n",
      "[2023-12-05 16:05:29,548] [INFO] [utils.py:802:see_memory_usage] After initializing ZeRO optimizer\n",
      "[2023-12-05 16:05:29,549] [INFO] [utils.py:803:see_memory_usage] MA 0.01 GB         Max_MA 0.25 GB         CA 2.62 GB         Max_CA 3 GB \n",
      "[2023-12-05 16:05:29,550] [INFO] [utils.py:810:see_memory_usage] CPU Virtual Memory:  used = 55.36 GB, percent = 25.6%\n",
      "[2023-12-05 16:05:29,551] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Final Optimizer = adamw\n",
      "[2023-12-05 16:05:29,551] [WARNING] [lr_schedules.py:759:__init__] total_num_steps 100 is less than warmup_num_steps 1000\n",
      "[2023-12-05 16:05:29,552] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed using configured LR scheduler = WarmupDecayLR\n",
      "[2023-12-05 16:05:29,552] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed LR Scheduler = <deepspeed.runtime.lr_schedules.WarmupDecayLR object at 0x7fb6523a1fc0>\n",
      "[2023-12-05 16:05:29,553] [INFO] [logging.py:96:log_dist] [Rank 0] step=0, skipped=0, lr=[0.0001], mom=[[0.9, 0.999]]\n",
      "[2023-12-05 16:05:29,554] [INFO] [config.py:974:print] DeepSpeedEngine configuration:\n",
      "[2023-12-05 16:05:29,554] [INFO] [config.py:978:print]   activation_checkpointing_config  {\n",
      "    \"partition_activations\": false, \n",
      "    \"contiguous_memory_optimization\": false, \n",
      "    \"cpu_checkpointing\": false, \n",
      "    \"number_checkpoints\": null, \n",
      "    \"synchronize_checkpoint_boundary\": false, \n",
      "    \"profile\": false\n",
      "}\n",
      "[2023-12-05 16:05:29,555] [INFO] [config.py:978:print]   aio_config ................... {'block_size': 1048576, 'queue_depth': 8, 'thread_count': 1, 'single_submit': False, 'overlap_events': True}\n",
      "[2023-12-05 16:05:29,555] [INFO] [config.py:978:print]   amp_enabled .................. False\n",
      "[2023-12-05 16:05:29,556] [INFO] [config.py:978:print]   amp_params ................... False\n",
      "[2023-12-05 16:05:29,556] [INFO] [config.py:978:print]   autotuning_config ............ {\n",
      "    \"enabled\": false, \n",
      "    \"start_step\": null, \n",
      "    \"end_step\": null, \n",
      "    \"metric_path\": null, \n",
      "    \"arg_mappings\": null, \n",
      "    \"metric\": \"throughput\", \n",
      "    \"model_info\": null, \n",
      "    \"results_dir\": \"autotuning_results\", \n",
      "    \"exps_dir\": \"autotuning_exps\", \n",
      "    \"overwrite\": true, \n",
      "    \"fast\": true, \n",
      "    \"start_profile_step\": 3, \n",
      "    \"end_profile_step\": 5, \n",
      "    \"tuner_type\": \"gridsearch\", \n",
      "    \"tuner_early_stopping\": 5, \n",
      "    \"tuner_num_trials\": 50, \n",
      "    \"model_info_path\": null, \n",
      "    \"mp_size\": 1, \n",
      "    \"max_train_batch_size\": null, \n",
      "    \"min_train_batch_size\": 1, \n",
      "    \"max_train_micro_batch_size_per_gpu\": 1.024000e+03, \n",
      "    \"min_train_micro_batch_size_per_gpu\": 1, \n",
      "    \"num_tuning_micro_batch_sizes\": 3\n",
      "}\n",
      "[2023-12-05 16:05:29,556] [INFO] [config.py:978:print]   bfloat16_enabled ............. True\n",
      "[2023-12-05 16:05:29,557] [INFO] [config.py:978:print]   checkpoint_parallel_write_pipeline  False\n",
      "[2023-12-05 16:05:29,557] [INFO] [config.py:978:print]   checkpoint_tag_validation_enabled  True\n",
      "[2023-12-05 16:05:29,557] [INFO] [config.py:978:print]   checkpoint_tag_validation_fail  False\n",
      "[2023-12-05 16:05:29,558] [INFO] [config.py:978:print]   comms_config ................. <deepspeed.comm.config.DeepSpeedCommsConfig object at 0x7fb427537a60>\n",
      "[2023-12-05 16:05:29,558] [INFO] [config.py:978:print]   communication_data_type ...... None\n",
      "[2023-12-05 16:05:29,558] [INFO] [config.py:978:print]   compression_config ........... {'weight_quantization': {'shared_parameters': {'enabled': False, 'quantizer_kernel': False, 'schedule_offset': 0, 'quantize_groups': 1, 'quantize_verbose': False, 'quantization_type': 'symmetric', 'quantize_weight_in_forward': False, 'rounding': 'nearest', 'fp16_mixed_quantize': False, 'quantize_change_ratio': 0.001}, 'different_groups': {}}, 'activation_quantization': {'shared_parameters': {'enabled': False, 'quantization_type': 'symmetric', 'range_calibration': 'dynamic', 'schedule_offset': 1000}, 'different_groups': {}}, 'sparse_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'row_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'head_pruning': {'shared_parameters': {'enabled': False, 'method': 'topk', 'schedule_offset': 1000}, 'different_groups': {}}, 'channel_pruning': {'shared_parameters': {'enabled': False, 'method': 'l1', 'schedule_offset': 1000}, 'different_groups': {}}, 'layer_reduction': {'enabled': False}}\n",
      "[2023-12-05 16:05:29,560] [INFO] [config.py:978:print]   curriculum_enabled_legacy .... False\n",
      "[2023-12-05 16:05:29,560] [INFO] [config.py:978:print]   curriculum_params_legacy ..... False\n",
      "[2023-12-05 16:05:29,560] [INFO] [config.py:978:print]   data_efficiency_config ....... {'enabled': False, 'seed': 1234, 'data_sampling': {'enabled': False, 'num_epochs': 1000, 'num_workers': 0, 'curriculum_learning': {'enabled': False}}, 'data_routing': {'enabled': False, 'random_ltd': {'enabled': False, 'layer_token_lr_schedule': {'enabled': False}}}}\n",
      "[2023-12-05 16:05:29,561] [INFO] [config.py:978:print]   data_efficiency_enabled ...... False\n",
      "[2023-12-05 16:05:29,561] [INFO] [config.py:978:print]   dataloader_drop_last ......... False\n",
      "[2023-12-05 16:05:29,561] [INFO] [config.py:978:print]   disable_allgather ............ False\n",
      "[2023-12-05 16:05:29,562] [INFO] [config.py:978:print]   dump_state ................... False\n",
      "[2023-12-05 16:05:29,562] [INFO] [config.py:978:print]   dynamic_loss_scale_args ...... None\n",
      "[2023-12-05 16:05:29,562] [INFO] [config.py:978:print]   eigenvalue_enabled ........... False\n",
      "[2023-12-05 16:05:29,563] [INFO] [config.py:978:print]   eigenvalue_gas_boundary_resolution  1\n",
      "[2023-12-05 16:05:29,564] [INFO] [config.py:978:print]   eigenvalue_layer_name ........ bert.encoder.layer\n",
      "[2023-12-05 16:05:29,564] [INFO] [config.py:978:print]   eigenvalue_layer_num ......... 0\n",
      "[2023-12-05 16:05:29,564] [INFO] [config.py:978:print]   eigenvalue_max_iter .......... 100\n",
      "[2023-12-05 16:05:29,565] [INFO] [config.py:978:print]   eigenvalue_stability ......... 1e-06\n",
      "[2023-12-05 16:05:29,565] [INFO] [config.py:978:print]   eigenvalue_tol ............... 0.01\n",
      "[2023-12-05 16:05:29,565] [INFO] [config.py:978:print]   eigenvalue_verbose ........... False\n",
      "[2023-12-05 16:05:29,566] [INFO] [config.py:978:print]   elasticity_enabled ........... False\n",
      "[2023-12-05 16:05:29,566] [INFO] [config.py:978:print]   flops_profiler_config ........ {\n",
      "    \"enabled\": false, \n",
      "    \"recompute_fwd_factor\": 0.0, \n",
      "    \"profile_step\": 1, \n",
      "    \"module_depth\": -1, \n",
      "    \"top_modules\": 1, \n",
      "    \"detailed\": true, \n",
      "    \"output_file\": null\n",
      "}\n",
      "[2023-12-05 16:05:29,567] [INFO] [config.py:978:print]   fp16_auto_cast ............... None\n",
      "[2023-12-05 16:05:29,567] [INFO] [config.py:978:print]   fp16_enabled ................. False\n",
      "[2023-12-05 16:05:29,567] [INFO] [config.py:978:print]   fp16_master_weights_and_gradients  False\n",
      "[2023-12-05 16:05:29,568] [INFO] [config.py:978:print]   global_rank .................. 0\n",
      "[2023-12-05 16:05:29,568] [INFO] [config.py:978:print]   grad_accum_dtype ............. None\n",
      "[2023-12-05 16:05:29,568] [INFO] [config.py:978:print]   gradient_accumulation_steps .. 1\n",
      "[2023-12-05 16:05:29,569] [INFO] [config.py:978:print]   gradient_clipping ............ 1.0\n",
      "[2023-12-05 16:05:29,569] [INFO] [config.py:978:print]   gradient_predivide_factor .... 1.0\n",
      "[2023-12-05 16:05:29,569] [INFO] [config.py:978:print]   hybrid_engine ................ enabled=False max_out_tokens=512 inference_tp_size=1 release_inference_cache=False pin_parameters=True tp_gather_partition_size=8\n",
      "[2023-12-05 16:05:29,570] [INFO] [config.py:978:print]   initial_dynamic_scale ........ 1\n",
      "[2023-12-05 16:05:29,571] [INFO] [config.py:978:print]   load_universal_checkpoint .... False\n",
      "[2023-12-05 16:05:29,571] [INFO] [config.py:978:print]   loss_scale ................... 1.0\n",
      "[2023-12-05 16:05:29,571] [INFO] [config.py:978:print]   memory_breakdown ............. False\n",
      "[2023-12-05 16:05:29,572] [INFO] [config.py:978:print]   mics_hierarchial_params_gather  False\n",
      "[2023-12-05 16:05:29,572] [INFO] [config.py:978:print]   mics_shard_size .............. -1\n",
      "[2023-12-05 16:05:29,573] [INFO] [config.py:978:print]   monitor_config ............... tensorboard=TensorBoardConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') wandb=WandbConfig(enabled=False, group=None, team=None, project='deepspeed') csv_monitor=CSVConfig(enabled=False, output_path='', job_name='DeepSpeedJobName') enabled=False\n",
      "[2023-12-05 16:05:29,573] [INFO] [config.py:978:print]   nebula_config ................ {\n",
      "    \"enabled\": false, \n",
      "    \"persistent_storage_path\": null, \n",
      "    \"persistent_time_interval\": 100, \n",
      "    \"num_of_version_in_retention\": 2, \n",
      "    \"enable_nebula_load\": true, \n",
      "    \"load_path\": null\n",
      "}\n",
      "[2023-12-05 16:05:29,573] [INFO] [config.py:978:print]   optimizer_legacy_fusion ...... False\n",
      "[2023-12-05 16:05:29,574] [INFO] [config.py:978:print]   optimizer_name ............... adamw\n",
      "[2023-12-05 16:05:29,574] [INFO] [config.py:978:print]   optimizer_params ............. {'lr': 0.0001, 'betas': [0.9, 0.999], 'eps': 1e-08, 'weight_decay': 0}\n",
      "[2023-12-05 16:05:29,574] [INFO] [config.py:978:print]   pipeline ..................... {'stages': 'auto', 'partition': 'best', 'seed_layers': False, 'activation_checkpoint_interval': 0, 'pipe_partitioned': True, 'grad_partitioned': True}\n",
      "[2023-12-05 16:05:29,575] [INFO] [config.py:978:print]   pld_enabled .................. False\n",
      "[2023-12-05 16:05:29,575] [INFO] [config.py:978:print]   pld_params ................... False\n",
      "[2023-12-05 16:05:29,575] [INFO] [config.py:978:print]   prescale_gradients ........... False\n",
      "[2023-12-05 16:05:29,576] [INFO] [config.py:978:print]   scheduler_name ............... WarmupDecayLR\n",
      "[2023-12-05 16:05:29,576] [INFO] [config.py:978:print]   scheduler_params ............. {'warmup_min_lr': 0, 'warmup_max_lr': 0.0001, 'warmup_num_steps': 1000, 'total_num_steps': 100}\n",
      "[2023-12-05 16:05:29,576] [INFO] [config.py:978:print]   seq_parallel_communication_data_type  torch.float32\n",
      "[2023-12-05 16:05:29,577] [INFO] [config.py:978:print]   sparse_attention ............. None\n",
      "[2023-12-05 16:05:29,577] [INFO] [config.py:978:print]   sparse_gradients_enabled ..... False\n",
      "[2023-12-05 16:05:29,577] [INFO] [config.py:978:print]   steps_per_print .............. inf\n",
      "[2023-12-05 16:05:29,578] [INFO] [config.py:978:print]   train_batch_size ............. 2\n",
      "[2023-12-05 16:05:29,578] [INFO] [config.py:978:print]   train_micro_batch_size_per_gpu  2\n",
      "[2023-12-05 16:05:29,578] [INFO] [config.py:978:print]   use_node_local_storage ....... False\n",
      "[2023-12-05 16:05:29,579] [INFO] [config.py:978:print]   wall_clock_breakdown ......... False\n",
      "[2023-12-05 16:05:29,579] [INFO] [config.py:978:print]   weight_quantization_config ... None\n",
      "[2023-12-05 16:05:29,579] [INFO] [config.py:978:print]   world_size ................... 1\n",
      "[2023-12-05 16:05:29,580] [INFO] [config.py:978:print]   zero_allow_untested_optimizer  False\n",
      "[2023-12-05 16:05:29,581] [INFO] [config.py:978:print]   zero_config .................. stage=3 contiguous_gradients=True reduce_scatter=True reduce_bucket_size=4194304 allgather_partitions=True allgather_bucket_size=500,000,000 overlap_comm=True load_from_fp32_weights=True elastic_checkpoint=False offload_param=DeepSpeedZeroOffloadParamConfig(device='cpu', nvme_path=None, buffer_count=5, buffer_size=100,000,000, max_in_cpu=1,000,000,000, pin_memory=True) offload_optimizer=DeepSpeedZeroOffloadOptimizerConfig(device='cpu', nvme_path=None, buffer_count=4, pin_memory=True, pipeline=False, pipeline_read=False, pipeline_write=False, fast_init=False, ratio=1.0) sub_group_size=100000000 cpu_offload_param=None cpu_offload_use_pin_memory=None cpu_offload=None prefetch_bucket_size=3774873 param_persistence_threshold=20480 model_persistence_threshold=sys.maxsize max_live_parameters=100000000 max_reuse_distance=100000000 gather_16bit_weights_on_model_save=True stage3_gather_fp16_weights_on_model_save=False ignore_unused_parameters=True legacy_stage1=False round_robin_gradients=False zero_hpz_partition_size=1 zero_quantized_weights=False zero_quantized_nontrainable_weights=False zero_quantized_gradients=False mics_shard_size=-1 mics_hierarchical_params_gather=False memory_efficient_linear=True pipeline_loading_checkpoint=False override_module_apply=True\n",
      "[2023-12-05 16:05:29,582] [INFO] [config.py:978:print]   zero_enabled ................. True\n",
      "[2023-12-05 16:05:29,582] [INFO] [config.py:978:print]   zero_force_ds_cpu_optimizer .. True\n",
      "[2023-12-05 16:05:29,582] [INFO] [config.py:978:print]   zero_optimization_stage ...... 3\n",
      "[2023-12-05 16:05:29,583] [INFO] [config.py:964:print_user_config]   json = {\n",
      "    \"comms_logger\": {\n",
      "        \"enabled\": true, \n",
      "        \"debug\": true\n",
      "    }, \n",
      "    \"fp16\": {\n",
      "        \"enabled\": false, \n",
      "        \"loss_scale\": 0, \n",
      "        \"loss_scale_window\": 1000, \n",
      "        \"initial_scale_power\": 16, \n",
      "        \"hysteresis\": 2, \n",
      "        \"min_loss_scale\": 1\n",
      "    }, \n",
      "    \"bf16\": {\n",
      "        \"enabled\": true\n",
      "    }, \n",
      "    \"optimizer\": {\n",
      "        \"type\": \"AdamW\", \n",
      "        \"params\": {\n",
      "            \"lr\": 0.0001, \n",
      "            \"betas\": [0.9, 0.999], \n",
      "            \"eps\": 1e-08, \n",
      "            \"weight_decay\": 0\n",
      "        }\n",
      "    }, \n",
      "    \"scheduler\": {\n",
      "        \"type\": \"WarmupDecayLR\", \n",
      "        \"params\": {\n",
      "            \"warmup_min_lr\": 0, \n",
      "            \"warmup_max_lr\": 0.0001, \n",
      "            \"warmup_num_steps\": 1000, \n",
      "            \"total_num_steps\": 100\n",
      "        }\n",
      "    }, \n",
      "    \"zero_optimization\": {\n",
      "        \"stage\": 3, \n",
      "        \"offload_optimizer\": {\n",
      "            \"device\": \"cpu\", \n",
      "            \"pin_memory\": true\n",
      "        }, \n",
      "        \"offload_param\": {\n",
      "            \"device\": \"cpu\", \n",
      "            \"pin_memory\": true\n",
      "        }, \n",
      "        \"overlap_comm\": true, \n",
      "        \"contiguous_gradients\": true, \n",
      "        \"sub_group_size\": 1.000000e+08, \n",
      "        \"reduce_bucket_size\": 4.194304e+06, \n",
      "        \"stage3_prefetch_bucket_size\": 3.774874e+06, \n",
      "        \"stage3_param_persistence_threshold\": 2.048000e+04, \n",
      "        \"stage3_max_live_parameters\": 1.000000e+08, \n",
      "        \"stage3_max_reuse_distance\": 1.000000e+08, \n",
      "        \"stage3_gather_16bit_weights_on_model_save\": true\n",
      "    }, \n",
      "    \"gradient_accumulation_steps\": 1, \n",
      "    \"gradient_clipping\": 1.0, \n",
      "    \"steps_per_print\": inf, \n",
      "    \"train_batch_size\": 2, \n",
      "    \"train_micro_batch_size_per_gpu\": 2, \n",
      "    \"wall_clock_breakdown\": false\n",
      "}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Attempting to resume from test-130m/checkpoint-5\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-12-05 16:05:29,590] [INFO] [torch_checkpoint_engine.py:27:load] [Torch] Loading checkpoint from test-130m/checkpoint-5/global_step5/zero_pp_rank_0_mp_rank_00_model_states.pt...\n",
      "[2023-12-05 16:05:29,601] [INFO] [torch_checkpoint_engine.py:29:load] [Torch] Loaded checkpoint from test-130m/checkpoint-5/global_step5/zero_pp_rank_0_mp_rank_00_model_states.pt.\n",
      "[2023-12-05 16:05:29,602] [INFO] [torch_checkpoint_engine.py:27:load] [Torch] Loading checkpoint from test-130m/checkpoint-5/global_step5/zero_pp_rank_0_mp_rank_00_model_states.pt...\n",
      "[2023-12-05 16:05:29,611] [INFO] [torch_checkpoint_engine.py:29:load] [Torch] Loaded checkpoint from test-130m/checkpoint-5/global_step5/zero_pp_rank_0_mp_rank_00_model_states.pt.\n",
      "[2023-12-05 16:05:29,621] [INFO] [torch_checkpoint_engine.py:27:load] [Torch] Loading checkpoint from test-130m/checkpoint-5/global_step5/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt...\n",
      "[2023-12-05 16:05:32,035] [INFO] [torch_checkpoint_engine.py:29:load] [Torch] Loaded checkpoint from test-130m/checkpoint-5/global_step5/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt.\n",
      "[2023-12-05 16:05:32,036] [INFO] [engine.py:2973:_get_all_zero_checkpoint_state_dicts] successfully read 1 ZeRO state_dicts for rank 0\n",
      "[2023-12-05 16:05:32,983] [INFO] [engine.py:2905:_load_zero_checkpoint] loading 1 zero partition checkpoints for rank 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "***** Running training *****\n",
      "  Num examples = 385,224\n",
      "  Num Epochs = 1\n",
      "  Instantaneous batch size per device = 2\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 2\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 100\n",
      "  Number of trainable parameters = 1,334,841,344\n",
      "  Continuing training from checkpoint, will skip to saved global_step\n",
      "  Continuing training from epoch 0\n",
      "  Continuing training from global step 5\n",
      "  Will skip the first 0 epochs then the first 5 batches in the first epoch.\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "104242\n",
      "261743\n",
      "235058\n",
      "178647\n",
      "Adam Optimizer #0 is created with AVX2 arithmetic capability.\n",
      "Config: alpha=0.000100, betas=(0.900000, 0.999000), weight_decay=0.000000, adam_w=1\n",
      "tensor(7.7500, device='cuda:0', dtype=torch.bfloat16,\n",
      "       grad_fn=<NllLossBackward0>) tensor([[-1.2109,  1.5234, -1.1094,  ..., -0.7695, -1.7109,  0.7188],\n",
      "        [-1.0000,  1.5156, -0.6523,  ..., -0.8203, -1.6875,  0.7656],\n",
      "        [-1.1797,  0.9805, -0.6641,  ..., -0.5586, -1.0234,  0.7148],\n",
      "        ...,\n",
      "        [-1.1016,  1.3828, -0.5547,  ..., -0.6758, -1.5547,  1.3438],\n",
      "        [-0.3125,  2.0625, -0.6211,  ..., -1.0938, -2.0312,  0.5898],\n",
      "        [-0.7656,  1.6641, -0.2109,  ..., -0.8594, -1.8359,  1.2578]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([   77,   201,    66,  ...,  4833,   521, 23351], device='cuda:0') torch.int64\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='11' max='100' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [ 11/100 00:10 < 03:46, 0.39 it/s, Epoch 0.00/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>7.750000</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>7.562500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>7.687500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>8.500000</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "225275\n",
      "290107\n",
      "tensor(7.5625, device='cuda:0', dtype=torch.bfloat16,\n",
      "       grad_fn=<NllLossBackward0>) tensor([[-1.0391,  1.9844, -0.9297,  ..., -1.6094, -1.9062,  1.2188],\n",
      "        [-1.1328,  1.6016, -0.9844,  ..., -0.8711, -1.8125,  0.9297],\n",
      "        [-0.5273,  1.7031, -0.9492,  ..., -1.3047, -1.6328,  0.7617],\n",
      "        ...,\n",
      "        [-0.7695,  1.7656, -1.0859,  ..., -1.0234, -2.4375,  1.2656],\n",
      "        [-0.4004,  1.9688, -1.2422,  ..., -0.9844, -2.1094,  1.3281],\n",
      "        [-0.5586,  2.1719, -0.7969,  ..., -0.6133, -2.1094,  0.8359]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([  267, 23724,  1206,  ...,   650, 29570,   628], device='cuda:0') torch.int64\n",
      "105315\n",
      "358768\n",
      "tensor(7.6875, device='cuda:0', dtype=torch.bfloat16,\n",
      "       grad_fn=<NllLossBackward0>) tensor([[-0.6523,  2.2500, -0.9883,  ..., -0.6523, -1.5625,  0.8555],\n",
      "        [-0.6836,  2.0781, -1.3047,  ..., -0.7031, -2.0938,  0.8594],\n",
      "        [-0.5273,  1.7031, -0.8633,  ..., -1.2500, -1.9922,  0.5117],\n",
      "        ...,\n",
      "        [-0.7070,  2.0781, -1.1094,  ..., -0.5898, -2.0781,  0.9805],\n",
      "        [-0.4824,  2.1406, -0.5938,  ..., -0.6016, -2.0156,  0.6523],\n",
      "        [-0.9219,  2.0469, -0.4980,  ..., -1.0391, -2.4688,  0.8477]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([2582, 2150,  492,  ...,  709,   17, 1999], device='cuda:0') torch.int64\n",
      "332421\n",
      "112040\n",
      "tensor(8.5000, device='cuda:0', dtype=torch.bfloat16,\n",
      "       grad_fn=<NllLossBackward0>) tensor([[-0.5352,  1.7344, -1.1562,  ..., -0.9570, -2.2188,  1.0156],\n",
      "        [-1.3750,  1.1172, -1.1016,  ..., -0.9727, -2.0781,  0.4316],\n",
      "        [ 0.3516,  1.7500, -0.5391,  ...,  0.0554, -2.2344,  0.5469],\n",
      "        ...,\n",
      "        [-1.0938,  1.7109, -1.1953,  ..., -0.7109, -2.0938,  0.4980],\n",
      "        [-0.3086,  2.1875, -1.0703,  ..., -0.3125, -1.8594,  0.6797],\n",
      "        [-0.8477,  1.5469, -0.5664,  ..., -0.7812, -2.5000,  0.4922]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([  201,    37, 12861,  ..., 13257,    17,  7103], device='cuda:0') torch.int64\n",
      "198053\n",
      "378413\n",
      "tensor(6.7188, device='cuda:0', dtype=torch.bfloat16,\n",
      "       grad_fn=<NllLossBackward0>) tensor([[-1.3594,  1.4531, -1.1328,  ..., -1.2969, -1.8047,  0.1641],\n",
      "        [-0.5781,  1.6328, -1.2969,  ..., -0.7266, -2.0156, -0.1045],\n",
      "        [-0.5820,  1.4219, -0.7305,  ..., -0.8281, -2.0000,  0.0535],\n",
      "        ...,\n",
      "        [-0.9844,  1.3828,  0.4180,  ..., -0.7461, -2.5938,  0.4395],\n",
      "        [-0.8477,  1.2031, -1.2109,  ..., -0.3848, -2.2812,  0.1973],\n",
      "        [-0.3848,  0.9766, -0.2500,  ..., -0.6758, -2.2656,  0.2715]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([ 1122,  8452, 16062,  ...,   201,    51,  7247], device='cuda:0') torch.int64\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Saving model checkpoint to test-130m/checkpoint-10\n",
      "Configuration saved in test-130m/checkpoint-10/config.json\n",
      "Removed shared tensor {'lm_head.weight'} while saving. This should be OK, but check by verifying that you don't receive any warning while reloading\n",
      "Model weights saved in test-130m/checkpoint-10/pytorch_model.bin\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[2023-12-05 16:06:05,475] [INFO] [logging.py:96:log_dist] [Rank 0] [Torch] Checkpoint global_step10 is about to be saved!\n",
      "[2023-12-05 16:06:05,484] [INFO] [logging.py:96:log_dist] [Rank 0] Saving model checkpoint: test-130m/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_model_states.pt\n",
      "[2023-12-05 16:06:05,484] [INFO] [torch_checkpoint_engine.py:21:save] [Torch] Saving test-130m/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_model_states.pt...\n",
      "[2023-12-05 16:06:05,504] [INFO] [torch_checkpoint_engine.py:23:save] [Torch] Saved test-130m/checkpoint-10/global_step10/zero_pp_rank_0_mp_rank_00_model_states.pt.\n",
      "[2023-12-05 16:06:05,506] [INFO] [torch_checkpoint_engine.py:21:save] [Torch] Saving test-130m/checkpoint-10/global_step10/bf16_zero_pp_rank_0_mp_rank_00_optim_states.pt...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/ubuntu/.local/lib/python3.10/site-packages/torch/nn/modules/module.py:1802: UserWarning: Positional args are being deprecated, use kwargs instead. Refer to https://pytorch.org/docs/master/generated/torch.nn.Module.html#torch.nn.Module.state_dict for details.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "trainer.train(resume_from_checkpoint = 'test-1.4b/checkpoint-5')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
